import json
import subprocess

# 从JSON文件加载配置
def load_config_from_json(file_path):
    with open(file_path, 'r') as f:
        return json.load(f)

# 替换可变参数（例如：${EPOCH} 替换成实际值）
def replace_parameters(config, **kwargs):
    for key, value in config.items():
        if isinstance(value, str):
            config[key] = value.format(**kwargs)
    return config

# 使用加载的配置运行命令
def run_training_with_config(config_file, **kwargs):
    # 加载配置
    config = load_config_from_json(config_file)
    
    # 替换可变参数
    config = replace_parameters(config, **kwargs)

    # 打印替换后的配置
    print("Final configuration after replacement:")
    for key, value in config.items():
        print(f"{key}: {value}")

    # 格式化命令
    command = [
        "accelerate", "launch",
        f"--config_file={config['config_file']}",
        f"--main_process_port={config['main_process_port']}",
        "train_model.py",
        f"--run_name={config['run_name']}",
        f"--pretrained_model_name_or_path={config['pretrained_model_name_or_path']}",
        f"--train_dataset_path={config['train_dataset_path']}",
        f"--eval_dataset_path={config['eval_dataset_path']}",
        f"--num_train_epochs={config['num_train_epochs']}",
        f"--learning_rate={config['learning_rate']}",
        f"--output_dir={config['output_dir']}",
        f"--eval_steps={config['eval_steps']}",
        f"--save_steps={config['save_steps']}",
        f"--max_length={config['max_length']}",
        f"--per_device_train_batch_size={config['per_device_train_batch_size']}",
        f"--per_device_eval_batch_size={config['per_device_eval_batch_size']}"
    ]
    
    # 运行命令
    subprocess.run(command)

# 调用函数执行训练
if __name__ == "__main__":
    #TODO
    config_file = 'train_config.json'  # 配置文件路径
    
    # 调用run_training_with_config并传入动态参数
    run_training_with_config(config_file, NAME="train", EPOCH=10, LR=5e-4)